//-------------------------------------------------------------------------------------
// ExportMesh.h
//
// Classes representing static and skinned meshes.  Also included is code to optimize
// mesh data and generate export-ready data from a non-indexed triangle list.
//  
// Advanced Technology Group (ATG)
// Copyright (C) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
//
// http://go.microsoft.com/fwlink/?LinkId=226208
//-------------------------------------------------------------------------------------
#pragma once

#include <d3d11.h>

// .SDKMESH uses Direct3D 9 decls, but only a subset of these is ever generated by the Content Exporter

enum D3DDECLUSAGE
{
    D3DDECLUSAGE_POSITION = 0,
    D3DDECLUSAGE_BLENDWEIGHT = 1,
    D3DDECLUSAGE_BLENDINDICES = 2,
    D3DDECLUSAGE_NORMAL = 3,
    D3DDECLUSAGE_TEXCOORD = 5,
    D3DDECLUSAGE_TANGENT = 6,
    D3DDECLUSAGE_BINORMAL = 7,
    D3DDECLUSAGE_COLOR = 10,
};

enum D3DDECLTYPE
{
    D3DDECLTYPE_FLOAT1 = 0,  // 1D float expanded to (value, 0., 0., 1.)
    D3DDECLTYPE_FLOAT2 = 1,  // 2D float expanded to (value, value, 0., 1.)
    D3DDECLTYPE_FLOAT3 = 2,  // 3D float expanded to (value, value, value, 1.)
    D3DDECLTYPE_FLOAT4 = 3,  // 4D float
    D3DDECLTYPE_D3DCOLOR = 4,  // 4D packed unsigned bytes mapped to 0. to 1. range
                                    // Input is in D3DCOLOR format (ARGB) expanded to (R, G, B, A)
                                    D3DDECLTYPE_UBYTE4 = 5,  // 4D unsigned byte
                                    D3DDECLTYPE_UBYTE4N = 8,  // Each of 4 bytes is normalized by dividing to 255.0
                                    D3DDECLTYPE_SHORT4N = 10,  // 4D signed short normalized (v[0]/32767.0,v[1]/32767.0,v[2]/32767.0,v[3]/32767.0)
                                    // Note: There is no equivalent to D3DDECLTYPE_DEC3N (14) as a DXGI_FORMAT
                                    D3DDECLTYPE_FLOAT16_2 = 15,  // Two 16-bit floating point values, expanded to (value, value, 0, 1)
                                    D3DDECLTYPE_FLOAT16_4 = 16,  // Four 16-bit floating point values

                                    D3DDECLTYPE_UNUSED = 17,  // When the type field in a decl is unused.

                                    // These are extensions for DXGI-based versions of Direct3D
                                    D3DDECLTYPE_DXGI_R10G10B10A2_UNORM = 32 + DXGI_FORMAT_R10G10B10A2_UNORM,
                                    D3DDECLTYPE_DXGI_R11G11B10_FLOAT = 32 + DXGI_FORMAT_R11G11B10_FLOAT,
                                    D3DDECLTYPE_DXGI_R8G8B8A8_SNORM = 32 + DXGI_FORMAT_R8G8B8A8_SNORM,
                                    D3DDECLTYPE_XBOX_R10G10B10_SNORM_A2_UNORM = 32 + 189,
};

#pragma pack(push,4)

struct D3DVERTEXELEMENT9
{
    WORD    Stream;     // Stream index
    WORD    Offset;     // Offset in the stream in bytes
    BYTE    Type;       // Data type
    BYTE    Method;     // Processing method
    BYTE    Usage;      // Semantics
    BYTE    UsageIndex; // Semantic index
};

#pragma pack(pop)

static_assert(sizeof(D3DVERTEXELEMENT9) == 8, "Direct3D9 Decl structure size incorrect");

namespace ATG
{

    class ExportVB
    {
    public:
        ExportVB()
            : m_uVertexCount(0),
            m_uVertexSizeBytes(0)
        {
        }
        ~ExportVB()
        {
        }

        void SetVertexSize(DWORD uByteCount) { m_uVertexSizeBytes = uByteCount; }
        DWORD GetVertexSize() const { return m_uVertexSizeBytes; }

        void SetVertexCount(size_t uVertexCount) { m_uVertexCount = uVertexCount; }
        size_t GetVertexCount() const { return m_uVertexCount; }

        void Allocate();

        uint8_t* GetVertex(size_t uIndex);
        const uint8_t* GetVertex(size_t uIndex) const;

        uint8_t* GetVertexData() { return m_pVertexData.get(); }
        const uint8_t* GetVertexData() const { return m_pVertexData.get(); }
        size_t GetVertexDataSize() const { return m_uVertexSizeBytes * m_uVertexCount; }

        void ByteSwap(const D3DVERTEXELEMENT9* pVertexElements, const size_t dwVertexElementCount);

    protected:
        DWORD                       m_uVertexSizeBytes;
        size_t                      m_uVertexCount;
        std::unique_ptr<uint8_t[]>  m_pVertexData;
    };

    class ExportIB
    {
    public:
        ExportIB()
            : m_uIndexCount(0),
            m_dwIndexSize(2)
        {
        }
        ~ExportIB()
        {
        }

        void SetIndexSize(DWORD dwIndexSize) { assert(dwIndexSize == 2 || dwIndexSize == 4); m_dwIndexSize = dwIndexSize; }
        DWORD GetIndexSize() const { return m_dwIndexSize; }

        void SetIndexCount(size_t uIndexCount) { m_uIndexCount = uIndexCount; }
        size_t GetIndexCount() const { return m_uIndexCount; }

        void Allocate();

        DWORD GetIndex(size_t uIndex) const
        {
            if (m_dwIndexSize == 2)
            {
                auto pIndexData16 = reinterpret_cast<const WORD*>(m_pIndexData.get());
                return pIndexData16[uIndex];
            }
            else
            {
                auto pIndexData32 = reinterpret_cast<const DWORD*>(m_pIndexData.get());
                return pIndexData32[uIndex];
            }
        }
        void SetIndex(size_t uIndex, DWORD dwData)
        {
            if (m_dwIndexSize == 2)
            {
                auto pIndexData16 = reinterpret_cast<WORD*>(m_pIndexData.get());
                pIndexData16[uIndex] = static_cast<WORD>(dwData);
            }
            else
            {
                auto pIndexData32 = reinterpret_cast<DWORD*>(m_pIndexData.get());
                pIndexData32[uIndex] = dwData;
            }
        }
        uint8_t* GetIndexData() { return m_pIndexData.get(); }
        const uint8_t* GetIndexData() const { return m_pIndexData.get(); }
        size_t GetIndexDataSize() const { return m_uIndexCount * m_dwIndexSize; }

        void ByteSwap();

    protected:
        DWORD                       m_dwIndexSize;
        size_t                      m_uIndexCount;
        std::unique_ptr<uint8_t[]>  m_pIndexData;
    };

    class ExportIBSubset :
        public ExportBase
    {
    public:
        enum PrimitiveType
        {
            TriangleList = 0,
            TriangleStrip,
            QuadList
        };
        ExportIBSubset()
            : m_uStartIndex(0),
            m_uIndexCount(0),
            m_PrimitiveType(TriangleList)
        {
        }
        void SetStartIndex(UINT uStartIndex) { m_uStartIndex = uStartIndex; }
        void IncrementIndexCount(UINT uSize) { m_uIndexCount += uSize; }
        void SetIndexCount(UINT uIndexCount) { m_uIndexCount = uIndexCount; }
        UINT GetStartIndex() const { return m_uStartIndex; }
        UINT GetIndexCount() const { return m_uIndexCount; }
        void SetPrimitiveType(PrimitiveType NewPT) { m_PrimitiveType = NewPT; }
        PrimitiveType GetPrimitiveType() const { return m_PrimitiveType; }
    protected:
        UINT            m_uStartIndex;
        UINT            m_uIndexCount;
        PrimitiveType   m_PrimitiveType;
    };

    class ExportMaterial;

    struct ExportMeshVertex
    {
    public:
        ExportMeshVertex() :
            DCCVertexIndex(0),
            pNextDuplicateVertex(nullptr)
        {
            Initialize();
        }
        void Initialize()
        {
            ZeroMemory(this, sizeof(ExportMeshVertex));
            BoneWeights.x = 1.0f;
        }
        UINT                            DCCVertexIndex;
        DirectX::XMFLOAT3               Position;
        DirectX::XMFLOAT3               Normal;
        DirectX::XMFLOAT3               SmoothNormal;
        DirectX::XMFLOAT3               Tangent;
        DirectX::XMFLOAT3               Binormal;
        DirectX::PackedVector::XMUBYTE4 BoneIndices;
        DirectX::XMFLOAT4               BoneWeights;
        DirectX::XMFLOAT4               TexCoords[8];
        DirectX::XMFLOAT4               Color;
        ExportMeshVertex* pNextDuplicateVertex;

        bool Equals(const ExportMeshVertex* pOtherVertex) const;
    };

    typedef std::vector< ExportMeshVertex* > ExportMeshVertexArray;

    struct ExportMeshTriangle
    {
    public:
        ExportMeshTriangle()
            : SubsetIndex(0),
            PolygonIndex(-1)
        {
        }
        void Initialize()
        {
            SubsetIndex = 0;
            Vertex[0].Initialize();
            Vertex[1].Initialize();
            Vertex[2].Initialize();
        }
        ExportMeshVertex    Vertex[3];
        INT                 SubsetIndex;
        INT                 PolygonIndex;
    };

    typedef std::vector< ExportMeshTriangle* > ExportMeshTriangleArray;

    class ExportMeshTriangleAllocator
    {
    public:
        ExportMeshTriangleAllocator()
            : m_uAllocatedCount(0),
            m_uTotalCount(0)
        {
        }
        ~ExportMeshTriangleAllocator()
        {
            Terminate();
        }
        void Initialize() { SetSizeHint(50000); }
        void Terminate();
        void SetSizeHint(UINT uAnticipatedSize);
        ExportMeshTriangle* GetNewTriangle();
        void ClearAllTriangles();
    private:
        struct AllocationBlock
        {
            ExportMeshTriangle* pTriangleArray;
            UINT m_uTriangleCount;
        };
        typedef std::list< AllocationBlock > AllocationBlockList;
        AllocationBlockList m_AllocationBlocks;
        UINT m_uTotalCount;
        UINT m_uAllocatedCount;
    };

    extern ExportMeshTriangleAllocator g_MeshTriangleAllocator;

    struct ExportVertexFormat
    {
    public:
        ExportVertexFormat()
            : m_bPosition(true),
            m_bNormal(true),
            m_bSkinData(false),
            m_bTangent(false),
            m_bBinormal(false),
            m_bVertexColor(true),
            m_uUVSetCount(1),
            m_uUVSetSize(2)
        {
        }
        bool        m_bPosition;
        bool        m_bNormal;
        bool        m_bTangent;
        bool        m_bBinormal;
        bool        m_bSkinData;
        bool        m_bVertexColor;
        UINT        m_uUVSetCount;
        UINT        m_uUVSetSize;
    };


    class ExportMeshBase :
        public ExportBase
    {
    public:
        ExportMeshBase(ExportString name);
        ~ExportMeshBase();

        enum MeshType
        {
            PolyMesh = 0
        };

        enum BoundsType
        {
            SphereBound = 0,
            AxisAlignedBoxBound = 1,
        };

        virtual MeshType GetMeshType() const = 0;

        void AddSubset(ExportIBSubset* pSubset) { m_vSubsets.push_back(pSubset); }
        size_t GetSubsetCount() const { return m_vSubsets.size(); }
        ExportIBSubset* GetSubset(size_t uIndex) { return m_vSubsets[uIndex]; }
        ExportIBSubset* FindSubset(const ExportString Name);

        DirectX::BoundingSphere& GetBoundingSphere() { return m_BoundingSphere; }
        DirectX::BoundingBox& GetBoundingAABB() { return m_BoundingAABB; }
        BoundsType GetSmallestBound() const { return m_SmallestBound; }

        virtual void AddInfluence(ExportString InfluenceName) { m_InfluenceNames.push_back(InfluenceName); }
        size_t GetInfluenceCount() const { return m_InfluenceNames.size(); }
        ExportString GetInfluence(size_t uIndex) const { return m_InfluenceNames[uIndex]; }

    protected:
        DirectX::BoundingSphere         m_BoundingSphere;
        DirectX::BoundingBox            m_BoundingAABB;
        BoundsType                      m_SmallestBound;
        std::vector< ExportIBSubset* >  m_vSubsets;
        std::vector< ExportString >     m_InfluenceNames;
    };


    class ExportSubDProcessMesh;

    class ExportMesh :
        public ExportMeshBase
    {
    public:
        enum OptimizationFlags
        {
            COMPRESS_VERTEX_DATA = 1,
            FLIP_TRIANGLES = 2,
            FORCE_SUBD_CONVERSION = 4,
            CLEAN_MESHES = 8,
            VCACHE_OPT = 16,
        };

        ExportMesh(ExportString name);
        ~ExportMesh();

        MeshType GetMeshType() const override { return ExportMeshBase::PolyMesh; }

        void SetVertexUVCount(UINT uCount) { m_VertexFormat.m_uUVSetCount = uCount; }
        void SetVertexUVDimension(UINT uDimension) { m_VertexFormat.m_uUVSetSize = uDimension; }
        void SetVertexNormalCount(UINT uCount);
        void SetVertexColorCount(UINT uCount) { m_VertexFormat.m_bVertexColor = (uCount > 0); }

        size_t GetVertexDeclElementCount() const { return m_VertexElements.size(); }
        const D3DVERTEXELEMENT9& GetVertexDeclElement(size_t uIndex) const { return m_VertexElements[uIndex]; }

        void AddRawTriangle(ExportMeshTriangle* pTriangle);
        void Optimize(DWORD dwFlags);
        void ByteSwap();

        ExportVB* GetVB() { return m_pVB.get(); }
        ExportIB* GetIB() { return m_pIB.get(); }

        ExportSubDProcessMesh* GetSubDMesh() { return m_pSubDMesh; }

        size_t GetTriangleCount() const { return m_TriangleToPolygonMapping.size(); }
        INT GetPolygonForTriangle(size_t dwTriangleIndex) const { return m_TriangleToPolygonMapping[dwTriangleIndex]; }

        void AddInfluence(ExportString InfluenceName) override { m_InfluenceNames.push_back(InfluenceName); m_VertexFormat.m_bSkinData = true; }

    protected:
        void BuildVertexBuffer(ExportMeshVertexArray& VertexArray, DWORD dwFlags);
        void ClearRawTriangles();
        void CleanMesh(bool breakBowTies);
        void ComputeVertexTangentSpaces();
        void ComputeAdjacency();
        void ComputeUVAtlas();
        void OptimizeVcache();
        void ComputeBoneSubsetGroups();
        void SortRawTrianglesBySubsetIndex();
        void ComputeBounds();

    protected:
        std::unique_ptr<ExportVB>                   m_pVB;
        std::unique_ptr<ExportIB>                   m_pIB;
        std::unique_ptr<DirectX::XMFLOAT3[]>        m_pVBPositions;
        std::unique_ptr<DirectX::XMFLOAT3[]>        m_pVBNormals;
        std::unique_ptr<DirectX::XMFLOAT2[]>        m_pVBTexCoords;
        std::unique_ptr<uint32_t[]>                 m_pAdjacency;
        std::unique_ptr<uint32_t[]>                 m_pAttributes;
        ExportMeshTriangleArray                     m_RawTriangles;
        std::vector< INT >                          m_TriangleToPolygonMapping;
        ExportVertexFormat                          m_VertexFormat;
        std::vector< D3DVERTEXELEMENT9 >            m_VertexElements;
        std::vector< D3D11_INPUT_ELEMENT_DESC  >    m_InputLayout;
        UINT                                        m_uDCCVertexCount;
        ExportSubDProcessMesh* m_pSubDMesh;
        bool                                        m_x2Bias;
    };

    class ExportMaterialSubsetBinding
    {
    public:
        ExportMaterialSubsetBinding()
        {
            ZeroMemory(this, sizeof(ExportMaterialSubsetBinding));
        }
        ExportString        SubsetName;
        ExportMaterial* pMaterial;
    };

    typedef std::vector< ExportMaterialSubsetBinding* > ExportMaterialSubsetBindingArray;

    class ExportModel
    {
    public:
        ExportModel(ExportMeshBase* pMesh)
            : m_pMesh(pMesh),
            m_bCastsShadows(true),
            m_bReceivesShadows(true)
        {
        }
        ~ExportModel();
        ExportMeshBase* GetMesh() { return m_pMesh; }
        bool SetSubsetBinding(ExportString SubsetName, ExportMaterial* pMaterial, bool bSkipValidation = false);
        size_t GetBindingCount() const { return m_vBindings.size(); }
        ExportMaterialSubsetBinding* GetBinding(size_t i) { return m_vBindings[i]; }

        bool IsShadowCaster() const { return m_bCastsShadows; }
        bool IsShadowReceiver() const { return m_bReceivesShadows; }
        void SetCastsShadows(bool bValue) { m_bCastsShadows = bValue; }
        void SetReceivesShadows(bool bValue) { m_bReceivesShadows = bValue; }

    protected:
        ExportMeshBase* m_pMesh;
        ExportMaterialSubsetBindingArray    m_vBindings;
        bool                                m_bCastsShadows;
        bool                                m_bReceivesShadows;
    };

};

